import os
import numpy as np
import tiktoken
from tqdm import trange,tqdm
import time
from src.model import APIModel
#from src.database import database
from src.utils import tokenCounter
from src.prompt import ROUGH_OUTLINE_PROMPT, MERGING_OUTLINE_PROMPT, SUBSECTION_OUTLINE_PROMPT, EDIT_FINAL_OUTLINE_PROMPT, SUGGESTION_OUTLINE_PROMPT, IMPROVE_FINAL_OUTLINE_PROMPT
import logging
import re

from pprint import pprint

class outlineWriter():
    
    def __init__(self, model:str, api_key:str, api_url:str, database) -> None:
        
        self.model, self.api_key, self.api_url = model, api_key, api_url 
        self.api_model = APIModel(self.model, self.api_key, self.api_url)

        self.db = database
        self.token_counter = tokenCounter()
        self.input_token_usage, self.output_token_usage = 0, 0

    def draft_outline(self, topic, reference_num=80, chunk_size=5000, section_num=5):
        references_ids = self.db.get_ids_from_query(topic, num=reference_num)
        references_infos = self.db.get_paper_info_from_ids(references_ids)

        references_titles = [r['title'] for r in references_infos]
        references_abs = [r['abs'] for r in references_infos]
        abs_chunks, titles_chunks = self.chunking(references_abs, references_titles, chunk_size=chunk_size)  # abs_chunks, titles_chunks 为二维数组，元素为str
        print('--- 论文检索完毕 ---')
        # 一级大纲
        outlines = self.generate_rough_outlines(topic=topic, papers_chunks=abs_chunks, titles_chunks=titles_chunks, section_num=section_num)
        print('--- 一级大纲生成完毕 ---')
        # merge 不同 chunk 的大纲 
        section_outline = self.merge_outlines(topic=topic, outlines=outlines)
        print('--- 一级大纲合并完毕 ---')
        subsection_outlines = self.generate_subsection_outlines(topic=topic, section_outline=section_outline, rag_num=10) # rag_num=30
        print('--- 二级大纲生成完毕 ---')
        merged_outline = self.process_outlines(section_outline, subsection_outlines)
        print('--- 大纲编辑完毕 ---')
        final_outline = self.edit_final_outline(merged_outline)
        print('--- 完整生成大纲完毕 ---')
        improve_outline = self.improvement_final_outline(final_outline)
        print('--- 大纲优化完毕 ---')
        return improve_outline

    def without_merging(self, topic, reference_num=600, chunk_size=30000, section_num=5):
        references_ids = self.db.get_ids_from_topic(topic = topic, num = reference_num, shuffle = False)
        references_infos = self.db.get_paper_info_from_ids(references_ids)

        references_titles = [r['title'] for r in references_infos]
        references_papers = [r['abs'] for r in references_infos]
        papers_chunks, titles_chunks = self.chunking(references_papers, references_titles, chunk_size=chunk_size)

        # 生成一级大纲 但不使用 merge，仅使用第一个 chunk 生成的大纲
        section_outline = self.generate_rough_outlines(topic=topic, papers_chunks = [papers_chunks[0]], titles_chunks = [titles_chunks[0]], section_num=section_num)[0]
        
        # generate subsection-level outline
        subsection_outlines = self.generate_subsection_outlines(topic=topic, section_outline=section_outline)
        
        final_outline = self.process_outlines(section_outline, subsection_outlines)

        return final_outline, section_outline, subsection_outlines

    def compute_price(self):
        return self.token_counter.compute_price(input_tokens=self.input_token_usage, output_tokens=self.output_token_usage, model=self.model)

    def generate_rough_outlines(self, topic, papers_chunks, titles_chunks, section_num=5):
        '''
        You wants to write a overall and comprehensive academic survey about "[TOPIC]".\n\
        You are provided with a list of papers related to the topic below:\n\
        ---
        [PAPER LIST]
        ---
        You need to draft a outline based on the given papers.
        The outline should contains a title and several sections.
        Each section follows with a brief sentence to describe what to write in this section.
        The outline is supposed to be comprehensive and contains [SECTION NUM] sections.

        Return in the format:
        <format>
        Title: [TITLE OF THE SURVEY]
        Section 1: [NAME OF SECTION 1]
        Description 1: [DESCRIPTION OF SENTCTION 1]

        Section 2: [NAME OF SECTION 2]
        Description 2: [DESCRIPTION OF SENTCTION 2]

        ...

        Section K: [NAME OF SECTION K]
        Description K: [DESCRIPTION OF SENTCTION K]
        </format>
        The outline:
        '''

        prompts = []
        for i in trange(len(papers_chunks)):
            titles = titles_chunks[i]
            papers = papers_chunks[i]
            paper_texts = '' 
            for i, t, p in zip(range(len(papers)), titles, papers):
                paper_texts += f'---\npaper_title: {t}\n\npaper_content:\n\n{p}\n'
            paper_texts+='---\n'
            prompt = self.__generate_prompt(ROUGH_OUTLINE_PROMPT, paras={'PAPER LIST': paper_texts, 'TOPIC': topic, 'SECTION NUM': str(section_num)})
            prompts.append(prompt)
        outlines = self.api_model.batch_chat(text_batch=prompts, temperature=1)
        return outlines
    
    def merge_outlines(self, topic, outlines):
        '''
        You are an expert in artificial intelligence who wants to write a overall survey about [TOPIC].\n\
        You are provided with a list of outlines as candidates below:\n\
        ---
        [OUTLINE LIST]
        ---
        Each outline contains a title and several sections.\n\
        Each section follows with a brief sentence to describe what to write in this section.\n\n\
        You need to generate a final outline based on these provided outlines.\n\
        Return in the format:
        <format>
        Title: [TITLE OF THE SURVEY]
        Section 1: [NAME OF SECTION 1]
        Description 1: [DESCRIPTION OF SENTCTION 1]

        Section 2: [NAME OF SECTION 2]
        Description 2: [DESCRIPTION OF SENTCTION 2]

        ...

        Section K: [NAME OF SECTION K]
        Description K: [DESCRIPTION OF SENTCTION K]
        </format>
        Only return the final outline without any other informations:
        '''
        outline_texts = '' 
        for i, o in zip(range(len(outlines)), outlines):
            outline_texts += f'---\noutline_id: {i}\n\noutline_content:\n\n{o}\n'
        outline_texts+='---\n'
        prompt = self.__generate_prompt(MERGING_OUTLINE_PROMPT, paras={'OUTLINE LIST': outline_texts, 'TOPIC':topic})

        outline = self.api_model.chat(prompt, temperature=1)
        return outline
    
    def generate_subsection_outlines(self, topic, section_outline, rag_num):
        '''
        You are an expert in artificial intelligence who wants to write a overall survey about [TOPIC].\n\
        You have created a overall outline below:\n\
        ---
        [OVERALL OUTLINE]
        ---
        The outline contains a title and several sections.\n\
        Each section follows with a brief sentence to describe what to write in this section.\n\n\
        <instruction>
        You need to enrich the section [SECTION NAME].
        The description of [SECTION NAME]: [SECTION DESCRIPTION]
        You need to generate the framwork containing several subsections based on the overall outlines.\n\
        Each subsection follows with a brief sentence to describe what to write in this subsection.
        These papers provided for references:
        ---
        [PAPER LIST]
        ---
        Return the outline in the format:
        <format>
        Subsection 1: [NAME OF SUBSECTION 1]
        Description 1: [DESCRIPTION OF SUBSENTCTION 1]

        Subsection 2: [NAME OF SUBSECTION 2]
        Description 2: [DESCRIPTION OF SUBSENTCTION 2]

        ...

        Subsection K: [NAME OF SUBSECTION K]
        Description K: [DESCRIPTION OF SUBSENTCTION K]
        </format>
        </instruction>
        Only return the outline without any other informations:
        '''
        survey_title, survey_sections, survey_section_descriptions = self.extract_title_sections_descriptions(section_outline)

        prompts = []

        for section_name, section_description in zip(survey_sections, survey_section_descriptions):
            references_ids = self.db.get_ids_from_query(section_description, num = rag_num, shuffle = True)
            references_infos = self.db.get_paper_info_from_ids(references_ids)

            references_titles = [r['title'] for r in references_infos]
            references_papers = [r['abs'] for r in references_infos]
            paper_texts = '' 
            for i, t, p in zip(range(len(references_papers)), references_titles, references_papers):
                paper_texts += f'---\npaper_title: {t}\n\npaper_content:\n\n{p}\n'
            paper_texts+='---\n'
            prompt = self.__generate_prompt(SUBSECTION_OUTLINE_PROMPT, paras={'OVERALL OUTLINE': section_outline,'SECTION NAME': section_name,\
                                                                          'SECTION DESCRIPTION':section_description,'TOPIC':topic,'PAPER LIST':paper_texts})
            prompts.append(prompt)

        sub_outlines = self.api_model.batch_chat(prompts, temperature=1)

        return sub_outlines
    
    def improvement_final_outline(self, outline):
        suggestion_prompt = self.__generate_prompt(SUGGESTION_OUTLINE_PROMPT, paras={'OUTLINE': outline})
        suggestion = self.api_model.chat(suggestion_prompt, temperature=1)
        suggestion = suggestion.replace('<Suggestion>','').replace('</Suggestion>','')

        improved_prompt = self.__generate_prompt(IMPROVE_FINAL_OUTLINE_PROMPT, paras={'OUTLINE': outline, 'SUGGESTION': suggestion})
        improved_outline = self.api_model.chat(improved_prompt, temperature=1)
        improved_outline = improved_outline.replace('<format>\n','').replace('</format>','')
        return improved_outline

    def edit_final_outline(self, outline):
        '''
        You are an expert in artificial intelligence who wants to write a overall survey about [TOPIC].\n\
        You have created a draft outline below:\n\
        ---
        [OVERALL OUTLINE]
        ---
        The outline contains a title and several sections.\n\
        Each section follows with a brief sentence to describe what to write in this section.\n\n\
        Under each section, there are several subsections.
        Each subsection also follows with a brief sentence of descripition.
        You need to modify the outline to make it both comprehensive and coherent with no repeated subsections.
        Return the final outline in the format:
        <format>
        # [TITLE OF SURVEY]

        ## [NAME OF SECTION 1]

        ### [NAME OF SUBSECTION 1]

        ### [NAME OF SUBSECTION 2]

        ...

        ### [NAME OF SUBSECTION L]
        
        ## [NAME OF SECTION 2]

        ...

        ## [NAME OF SECTION K]

        ...
        </format>
        Only return the final outline without any other informations:
        '''

        prompt = self.__generate_prompt(EDIT_FINAL_OUTLINE_PROMPT, paras={'OVERALL OUTLINE': outline})
        # self.input_token_usage += self.token_counter.num_tokens_from_string(prompt)
        outline = self.api_model.chat(prompt, temperature=1)
        # self.output_token_usage += self.token_counter.num_tokens_from_string(outline)
        return outline.replace('<format>\n','').replace('</format>','')
 
    def __generate_prompt(self, template, paras):
        # 生成提示词
        prompt = template
        for k in paras.keys():
            prompt = prompt.replace(f'[{k}]', paras[k])
        return prompt
    
    def extract_title_sections_descriptions(self, outline):
        title = outline.split('Title: ')[1].split('\n')[0]
        # 初始化章节和描述的列表
        sections, descriptions = [], []
        # 遍历可能的章节和描述，最多遍历 100 个章节
        for i in range(100):
            # 如果大纲中存在对应的章节，提取章节名和描述
            if f'Section {i+1}' in outline:
                sections.append(outline.split(f'Section {i+1}: ')[1].split('\n')[0])
                descriptions.append(outline.split(f'Description {i+1}: ')[1].split('\n')[0])
        return title, sections, descriptions
    
    def extract_subsections_subdescriptions(self, outline):
        """
        从大纲中提取小节和描述
        Args:
            outline: 大纲文本
        Returns:
            tuple: (小节列表, 描述列表)
        """
        subsections, subdescriptions = [], []
        
        # 输入验证
        if not outline or not isinstance(outline, str):
            logging.warning("输入的outline为空或格式不正确")
            return subsections, subdescriptions
            
        try:
            subsection_pattern = r'Subsection\s*(\d+)\s*:\s*([^\n]+)'
            description_pattern = r'Description\s*(\d+)\s*:\s*([^\n]+)'
            
            # 找到所有匹配项
            subsection_matches = re.findall(subsection_pattern, outline)
            description_matches = re.findall(description_pattern, outline)
            
            if not subsection_matches:
                logging.warning(f"未找到符合格式的小节内容: {outline[:100]}...")
                return [], []
            
            # 按序号排序
            subsection_dict = {int(num): content.strip() for num, content in subsection_matches}
            description_dict = {int(num): content.strip() for num, content in description_matches}
            
            # 确保序号连续且对应
            max_num = max(subsection_dict.keys()) if subsection_dict else 0
            for i in range(1, max_num + 1):
                if i in subsection_dict:
                    subsections.append(subsection_dict[i])
                    # 如果没有对应的描述，使用空字符串
                    subdescriptions.append(description_dict.get(i, ""))
            
            logging.info(f"成功提取 {len(subsections)} 个小节")
            
        except Exception as e:
            logging.error(f"提取小节时发生错误: {str(e)}")
            return [], []
            
        return subsections, subdescriptions
    
    def chunking(self, papers, titles, chunk_size=14000):
        paper_chunks, title_chunks = [], []
        # 计算输入论文的总 token 数量
        total_length = self.token_counter.num_tokens_from_list_string(papers)

        # 根据总 token 数量和设定的块大小，计算需要将文本分为多少个块
        num_of_chunks = int(total_length / chunk_size) + 1

        # 计算每个块的平均长度（token 数量），加 1 用于向上取整，避免 token 不足
        avg_len = int(total_length / num_of_chunks) + 1

        # split_points 用于存储每个分割点的位置
        split_points = []
        l = 0
        # 遍历所有论文内容，找到分割点
        for j in range(len(papers)):
            # 计算当前论文内容的 token 数量并累加
            l += self.token_counter.num_tokens_from_string(papers[j])
            # 如果累加的 token 数量大于每个块的平均长度，将当前位置作为分割点
            if l > avg_len:
                l = 0  # 重置累加器
                split_points.append(j)  # 将当前索引添加到分割点列表
                continue
        # 按照找到的分割点将论文内容和标题进行分块
        start = 0
        for point in split_points:
            paper_chunks.append(papers[start:point])  # 将 start 到分割点之间的内容作为一个块
            title_chunks.append(titles[start:point])  # 将对应的标题也进行分块
            start = point  # 更新起始位置
        
        # 将剩余的内容作为最后一个块
        paper_chunks.append(papers[start:])
        title_chunks.append(titles[start:])
        
        # 返回分块后的论文内容和标题
        return paper_chunks, title_chunks
       
    def process_outlines(self, section_outline, sub_outlines):
        """
        处理大纲和子大纲，生成最终的大纲文档
        Args:
            section_outline: 章节大纲
            sub_outlines: 子大纲列表
        Returns:
            str: 处理后的大纲文档
        """
        # 输入验证
        if not section_outline or not sub_outlines:
            logging.error("section_outline 或 sub_outlines 为空")
            return ""
            
        if not isinstance(sub_outlines, list):
            logging.error("sub_outlines 格式不正确")
            return ""
            
        try:
            # 初始化结果字符串
            res = ''
            # 提取调查的标题、章节和章节描述
            survey_title, survey_sections, survey_section_descriptions = self.extract_title_sections_descriptions(outline=section_outline)
            
            if not survey_sections:
                logging.error("未能提取到有效的章节信息")
                return ""
                
            # 将标题添加到结果字符串中
            res += f'# {survey_title}\n\n'
            
            # 遍历每个章节，将章节标题和描述添加到结果字符串中
            for i in range(len(survey_sections)):
                if i >= len(sub_outlines):
                    logging.warning(f"章节 {i+1} 缺少子大纲")
                    continue
                    
                section = survey_sections[i]
                res += f'## {i+1} {section}\nDescription: {survey_section_descriptions[i]}\n\n'
                
                try:
                    # 提取子章节和子描述，并将其添加到结果字符串中
                    subsections, subsection_descriptions = self.extract_subsections_subdescriptions(sub_outlines[i])
                    
                    if not subsections:
                        logging.warning(f"章节 {i+1} ({section}) 未找到有效的小节")
                        continue
                        
                    for j in range(len(subsections)):
                        subsection = subsections[j]
                        description = subsection_descriptions[j] if j < len(subsection_descriptions) else ""
                        res += f'### {i+1}.{j+1} {subsection}\nDescription: {description}\n\n'
                        
                except Exception as e:
                    logging.error(f"处理章节 {i+1} ({section}) 的小节时发生错误: {str(e)}")
                    continue
            
            if not res.strip():
                logging.error("生成的大纲内容为空")
                return ""
            # 返回处理后的大纲文档
            return res
            
        except Exception as e:
            logging.error(f"处理大纲时发生错误: {str(e)}")
            return ""

